#include <stdio.h>
#include <vector>

/**
 * @brief 使用C++定义并查集类
 * 
 */
class UnionSet {
public:
    UnionSet(const int n) {
        _n = n;
        for (int i = 0; i < _n; ++i) {
            _set.push_back(i);
            _height.push_back(0);   //! 当节点的集是自己的时候高度为0
        }
    }

    // 合并一个组关系
    void joint(int x, int y) {
        int set_x = find(x);
        int set_y = find(y);

        if (set_x == set_y)
            return;
        
        // 需要判断集的高度，而不是节点本身的高度
        if (_height[set_x] == _height[set_y]) {
            _set[set_x] = set_y;
            _height[set_y]++;
        } else {
            // 不一样高的时候，我们要将矮的合并到高的上去
            // 就是让矮的集指向高的集
            if (_height[set_x] < _height[set_y]) {
                _set[set_x] = set_y;
            } else {
                _set[set_y] = set_x;
            }
        }
    }

    // 查询一个元素所属的集
    int find(int x) {
        if (x != _set[x]) _set[x] = find(_set[x]);
        return _set[x];
    }

    int size() const {
        return _n;
    }

    int height(int x) const {
        return _height[x];
    }

private:
    std::vector<int> _set;
    std::vector<int> _height;
    int _n;
};

//
int main(int argc, char* argv[]) {
    UnionSet us(10);

    us.joint(1, 2);
    us.joint(2, 3);
    us.joint(3, 4);
    us.joint(5, 6);  //
    us.joint(6, 7);
    us.joint(7, 8);

    for (int i = 0; i < us.size(); i++) {
        int sset = us.find(i);
        printf("set[%d] = %d, height[%d] = %d\n", i, sset , i, us.height(sset));
    }

    printf("1 and 2 belong to the same set: %d\n", us.find(1) == us.find(2));
    printf("3 and 4 belong to the same set: %d\n", us.find(3) == us.find(4));
    printf("4 and 8 belong to the same set: %d\n", us.find(4) == us.find(8));
    return 0;
}